"""
知识库数据模式
"""

from typing import Optional, List, Dict, Any
from datetime import datetime
from pydantic import BaseModel


class KnowledgeBaseBase(BaseModel):
    """知识库基础模式"""
    name: str
    description: Optional[str] = None
    embedding_model: str = "text-embedding-ada-002"
    chunk_size: int = 1000
    chunk_overlap: int = 200


class KnowledgeBaseCreate(KnowledgeBaseBase):
    """创建知识库模式"""
    pass


class KnowledgeBaseUpdate(BaseModel):
    """更新知识库模式"""
    name: Optional[str] = None
    description: Optional[str] = None
    embedding_model: Optional[str] = None
    chunk_size: Optional[int] = None
    chunk_overlap: Optional[int] = None
    is_active: Optional[bool] = None


class KnowledgeBaseResponse(KnowledgeBaseBase):
    """知识库响应模式"""
    id: int
    collection_name: str
    is_active: bool
    document_count: int
    total_chunks: int
    created_at: datetime
    updated_at: datetime
    
    class Config:
        from_attributes = True


class DocumentBase(BaseModel):
    """文档基础模式"""
    filename: str
    original_filename: str
    file_path: str
    file_size: int
    file_type: str
    content: Optional[str] = None
    metadata: Dict[str, Any] = {}


class DocumentCreate(DocumentBase):
    """创建文档模式"""
    knowledge_base_id: int


class DocumentResponse(DocumentBase):
    """文档响应模式"""
    id: int
    knowledge_base_id: int
    chunk_count: int
    is_processed: bool
    processing_status: str
    error_message: Optional[str] = None
    created_at: datetime
    updated_at: datetime
    
    class Config:
        from_attributes = True
    
    @classmethod
    def from_orm(cls, obj):
        """自定义ORM转换，处理metadata字段"""
        # 处理metadata字段，确保它是字典类型
        metadata = obj.document_metadata
        if metadata is None:
            metadata = {}
        elif not isinstance(metadata, dict):
            metadata = {}
        
        data = {
            'id': obj.id,
            'knowledge_base_id': obj.knowledge_base_id,
            'filename': obj.filename,
            'original_filename': obj.original_filename,
            'file_path': obj.file_path,
            'file_size': obj.file_size,
            'file_type': obj.file_type,
            'content': obj.content,
            'metadata': metadata,
            'chunk_count': obj.chunk_count,
            'is_processed': obj.is_processed,
            'processing_status': obj.processing_status,
            'error_message': obj.error_message,
            'created_at': obj.created_at,
            'updated_at': obj.updated_at
        }
        return cls(**data)


class KnowledgeBaseWithDocuments(KnowledgeBaseResponse):
    """包含文档的知识库响应模式"""
    documents: List[DocumentResponse] = []


class DocumentUploadRequest(BaseModel):
    """文档上传请求模式"""
    knowledge_base_id: int
    filename: str
    content: str  # Base64编码的文件内容


class DocumentUploadResponse(BaseModel):
    """文档上传响应模式"""
    document_id: int
    filename: str
    file_size: int
    processing_status: str
    message: str


class KnowledgeBaseStats(BaseModel):
    """知识库统计模式"""
    id: int
    name: str
    description: Optional[str]
    document_count: int
    total_chunks: int
    embedding_model: str
    chunk_size: int
    chunk_overlap: int
    is_active: bool
    created_at: datetime
    updated_at: datetime
    documents: List[Dict[str, Any]] = []
